import tkinter as tk
from tkinter import ttk
from transformers import pipeline, AutoModelWithLMHead, AutoTokenizer
import warnings
import threading
import pyperclip


def app():
    root = main()
    root.mainloop()


def main():
    root = tk.Tk()
    root.geometry("187x120")
    get_window_width_height(root)  # 设置窗口在屏幕右侧靠中间位置
    root.title("本地翻译器")
    root.wm_attributes("-topmost", True)  # 窗口置顶
    root.wm_attributes("-toolwindow", True)  # 设为工具窗口，取消最大最小化

    # 选择框
    box = ttk.Combobox(root, width=20, state="readonly")
    box['value'] = ['中文转英文', '英文转中文']
    box.grid(row=1, column=1, columnspan=2, padx=2, pady=2)

    # 选择绑定事件函数
    def get_bind(event):
        global box_index
        if box.get() == "中文转英文":
            box_index = 0
        elif box.get() == "英文转中文":
            box_index = 1

    box.bind("<<ComboboxSelected>>", get_bind)
    box.current(0)  # 默认显示下标0的内容

    # 输入文本框
    entry = tk.Entry(root, width=15)
    entry.grid(row=2, column=1, padx=2, pady=2)

    # 翻译按钮
    tran_btn = tk.Button(root, text="翻译", width=5,
                         command=lambda: threading.Thread(target=tran_btn_click, args=(entry.get(), label)).start())
    tran_btn.grid(row=2, column=2, padx=2, pady=2)

    # 翻译内容文本框
    label = tk.Label(root, text="翻译的内容", anchor='w', width=15)
    label.grid(row=3, column=1, padx=2, pady=2)

    # 复制翻译内容按钮
    def copy_btn_click():
        copy_text = label.cget("text")  # 获取label的内容
        pyperclip.copy(copy_text)  # 将内容写入剪切板

    copy_btn = tk.Button(root, text="复制", width=5, command=copy_btn_click)
    copy_btn.grid(row=3, column=2, padx=2, pady=2)

    return root


def tran_btn_click(text, label):
    global box_index, tran0, tran1
    result = ''
    if box_index == 0:
        result = tran0(text)[0]['translation_text']  # 获取翻译内容
    elif box_index == 1:
        result = tran1(text)[0]['translation_text']  # 获取翻译内容

    if result != "":
        tran_text = result.rstrip('.')  # 去除无关字符
        label.config(text=tran_text)  # 将翻译后的内容配置到label控件里


def init_translation():
    model = AutoModelWithLMHead.from_pretrained('opus-mt-zh-en')  # 从预训练模型中加载模型
    tokenizer = AutoTokenizer.from_pretrained('opus-mt-zh-en')  # 从预训练模型的标记器中加载标记器
    translation1 = pipeline('translation_zh_to_en', model=model, tokenizer=tokenizer)  # 创建一个文本翻译的管道

    model = AutoModelWithLMHead.from_pretrained('opus-mt-en-zh')  # 从预训练模型中加载模型
    tokenizer = AutoTokenizer.from_pretrained('opus-mt-en-zh')  # 从预训练模型的标记器中加载标记器
    translation2 = pipeline('translation_en_to_zh', model=model, tokenizer=tokenizer)  # 创建一个文本翻译的管道

    return translation1, translation2


def get_window_width_height(window):
    screen_width = window.winfo_screenwidth()
    screen_height = window.winfo_screenheight()
    x = screen_width - 200
    y = int(screen_height / 2)
    window.geometry(f"+{x}+{y}")


if __name__ == '__main__':
    warnings.filterwarnings('ignore')  # 隐藏ignore信息
    box_index = 0  # 确定是哪种模式
    tran0, tran1 = init_translation()  # 初始化翻译模型
    app()
